import json
from args import *


def get_area(area_str):
    for area in areas:
        if area in area_str:
            return area
    return '其他'


def get_acc(model_predict_path: str, gt_path: str) -> float:
    with open(model_predict_path, "r") as f:
        data = json.load(f)

    for d in data:
        area_str = d['area']
        d['area'] = get_area(area_str)

    with open(gt_path, 'r') as f:
        data_gt = json.load(f)

    for i in range(len(data)):
        data[i]['gt_area'] = [i[0] for i in data_gt[i]['area']]

    # with open('data_compare.json', 'w') as f:
    #     json.dump(data, f, ensure_ascii=False, indent=4)
    total_len = len(data)
    correct = 0

    for i in range(len(data)):
        d = data[i]
        predict = d['area']
        gt = d['gt_area']
        if predict in gt:
            correct += 1

    return correct / total_len


if __name__ == '__main__':
    model_predict_path = ["area_chatglm3-6b.json", 'area_qwen1.5-7b-chat.json', 'area_qwen1.5-14b-chat.json']
    gt_path = 'data.json'
    for model_path in model_predict_path:
        acc = get_acc(model_path, gt_path)
        print(f'model_name: {model_path} acc: {acc:.3f}')
